import pandas as pd

# 读取total.csv文件
df = pd.read_csv('expendData/total_extend.csv')

# 假设最后一列是标签列
label_column = df.columns[-1]

# 创建一个空的DataFrame来存储选中的数据
selected_data = pd.DataFrame()

# 初始化一个空数组来存储原始行号
original_indices = []

# 为每个标签（除了BENIGN）选取前100条数据
for label in df[label_column].unique():
    if label == 'BENIGN':
        # BENIGN标签取500条数据
        selected_benign = df[df[label_column] == 'BENIGN'].head(500)
        selected_data = pd.concat([selected_data, selected_benign], ignore_index=True)
        # 添加BENIGN的行号到数组中
        original_indices.extend(selected_benign.index.tolist())
    else:
        # 其他标签取前100条数据
        selected_label = df[df[label_column] == label].head(100)
        selected_data = pd.concat([selected_data, selected_label], ignore_index=True)
        # 添加其他标签的行号到数组中
        original_indices.extend(selected_label.index.tolist())

    # 在DataFrame中添加一个列来保存原始行号（如果需要）
#selected_data['Original_Index'] = original_indices

# 保存选中的数据到新的CSV文件，不包括pandas自动添加的索引
selected_data.to_csv('total_small.csv', index=False)

# 如果你只需要行号数组，而不需要添加到DataFrame中，你可以直接使用original_indices数组
print(original_indices)
df = pd.read_csv('Normalization_total_extend.csv')

# 检查original_indices中的值是否都是整数，并且都在DataFrame的行数范围内
# 如果不是，你可能需要处理这些异常值
assert all(
    isinstance(idx, int) and 0 <= idx < len(df) for idx in original_indices), "Invalid indices in original_indices"

# 使用iloc根据行号选择行（注意iloc使用基于0的索引）
selected_rows = df.iloc[original_indices]

# 将选定的行保存到新的CSV文件中
selected_rows.to_csv('Normalization_total_small.csv', index=False)
print("ok")